from flask import Flask, request,render_template, Response, jsonify
import pyrealsense2 as rs
import numpy as np
import cv2
import threading
import time
from ultralytics import YOLO

model = YOLO("/home/rm/Orsen/New_Retail/best.pt")
kele = cv2.imread("static/kele.jpg")
results = model.predict(kele, conf=0.5, imgsz=640)

app = Flask(__name__)

# 全局变量存储最新帧
latest_color_frame = None
latest_depth_frame = None
intr = None
frame_lock = threading.Lock()

# 初始化RealSense管道
pipeline = rs.pipeline()
config = rs.config()
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)

def camera_thread():
    """持续获取相机画面的线程"""
    global latest_color_frame
    global latest_depth_frame
    global intr
    global depth_scale
    # pipeline.start(config)
    profile = pipeline.start(config)
    # 获取对齐工具
    align_to = rs.stream.color
    align = rs.align(align_to)
    depth_sensor = profile.get_device().first_depth_sensor()
    depth_scale = depth_sensor.get_depth_scale()
    print(f"深度比例: {depth_scale}")
    
    # 获取内参
    intr = profile.get_stream(rs.stream.color).as_video_stream_profile().get_intrinsics()
    # camera_intrinsic = o3d.camera.PinholeCameraIntrinsic(
    #     intr.width, intr.height, intr.fx, intr.fy, intr.ppx, intr.ppy
    # )

    try:
        while True:
            frames = pipeline.wait_for_frames()
            color_frame = frames.get_color_frame()
            if not color_frame:
                continue
            # 获取对齐帧
            aligned_frames = align.process(frames)
            
            # 检查帧有效性
            color_frame = aligned_frames.get_color_frame()
            depth_frame = aligned_frames.get_depth_frame()
            # if not color_frame or not depth_frame:
            #     continue
            #     raise RuntimeError("帧获取失败")
            
            # print(f"深度中值: {depth_frame.get_distance(320, 240):.2f}米")
            
            # 转换为Open3D格式
            color_image = np.asanyarray(color_frame.get_data())
            depth_image = np.asanyarray(depth_frame.get_data())
            
            # 检查图像数据
            # if np.max(depth_image) == 0:
            #     raise ValueError("深度图为空！")
                
            # # 转换为OpenCV格式
            # frame = np.asanyarray(color_frame.get_data())
            
            # 更新全局帧
            with frame_lock:
                latest_color_frame = color_image.copy()
                latest_depth_frame = depth_image.copy()
                
            time.sleep(0.033)  # ~30fps
            
    finally:
        pipeline.stop()

def pixel2point3d(px, py, depth_mm, camera_param):
    '''将像素坐标转换为相机坐标系下的三维坐标
    @px: 像素坐标 X轴坐标
    @py: 像素坐标 Y轴坐标
    @depth_mm: 深度值
    @camera_param: 相机参数
    '''
    # 获取彩色相机内参
    fx = camera_param.fx
    fy = camera_param.fy
    cx = camera_param.ppx
    cy = camera_param.ppy

    # 计算相机坐标系下的三维坐标
    cam_z = depth_mm
    cam_x = (px - cx) * cam_z / fx
    cam_y = (py - cy) * cam_z / fy

    return [cam_x, cam_y, cam_z]

@app.route('/')
def index():
    """显示实时画面的网页"""
    return render_template('index.html')

@app.route('/getIntr')
def get_intr():
    return jsonify({
        "cx":intr.ppx,"cy":intr.ppy,"fx":intr.fx,"fy":intr.fy
        })

def find_nearest_rectangle(rectangles):
    """
    从矩形框集合中找出深度最小的矩形
    
    参数:
        rectangles (list): 矩形框集合，每个矩形格式为 [x1, y1, x2, y2, depth]
        
    返回:
        list: 面积最大的矩形 [x1, y1, x2, y2, depth]
        或 None (如果输入为空)
    """
    if not rectangles:
        return None
    
    min_depth = 500
    nearest_rect = None
    
    for rect in rectangles:
        # 检查矩形格式是否有效
        if len(rect) != 5:
            raise ValueError("每个矩形必须包含4个坐标值 [x1, y1, x2, y2, depth]")
        
        x1, y1, x2, y2, depth = rect
        
        if depth < min_depth:
            min_depth = depth
            nearest_rect = rect

    return nearest_rect

def find_largest_rectangle(rectangles):
    """
    从矩形框集合中找出面积最大的矩形
    
    参数:
        rectangles (list): 矩形框集合，每个矩形格式为 [x1, y1, x2, y2, depth]
        
    返回:
        list: 面积最大的矩形 [x1, y1, x2, y2, depth]
        或 None (如果输入为空)
    """
    if not rectangles:
        return None
    
    max_area = -1
    largest_rect = None
    
    for rect in rectangles:
        # 检查矩形格式是否有效
        if len(rect) != 5:
            raise ValueError("每个矩形必须包含4个坐标值 [x1, y1, x2, y2, depth]")
        
        x1, y1, x2, y2, depth = rect
        
        # 计算当前矩形面积 (宽度 * 高度)
        width = abs(x2 - x1)
        height = abs(y2 - y1)
        area = width * height
        
        # 更新最大面积矩形
        if area > max_area:
            max_area = area
            largest_rect = rect
    
    return largest_rect

def draw_top1(latest_color_frame,box):
    x1 = box[0]
    y1 = box[1]
    x2 = box[2]
    y2 = box[3]
    target_c_x = int((x2-x1)/2+x1)
    target_c_y = int((y2-y1)/2+y1)
    center_coordinates = (target_c_x, target_c_y)
    radius = 5  
    color = (0, 255, 0)  # B=0, G=255, R=0  
    image = cv2.circle(latest_color_frame, center_coordinates, radius, color, -1)
    
    top_left = (box[0], box[1])
    bottom_right = (box[2], box[3])
    color = (0, 0, 255)  # BGR格式的颜色，绿色
    cv2.rectangle(latest_color_frame, top_left, bottom_right, color, thickness=2)

def filter_invalid_data(data_list, min_val=200, max_val=500):
    """
    过滤掉无效数据（0值或超出范围的数据）
    
    参数:
        data_list (list): 原始数据列表
        min_val (int): 有效最小值（默认200，过滤0）
        max_val (int): 有效最大值（默认500）
        
    返回:
        list: 过滤后的有效数据集合（保持原始顺序）
    """
    new_list = []
    for item in data_list:
        if item > min_val and item < max_val:
            new_list.append(item)
    return new_list    


@app.route('/findSku')
def find_sku():
    sku_id = request.args.get('sku_id', None)
    sku_scale = request.args.get('sku_scale', None)
    print(sku_id,sku_scale)

    with frame_lock:
        if latest_color_frame is None:
            return jsonify({"error": "相机未就绪"}), 500
            
        # 保存临时图片
        timestamp = int(time.time())
        filename = f"static/capture_{timestamp}.jpg"
        cv2.imwrite(filename, latest_color_frame)
        results = model.predict(latest_color_frame, conf=0.5, imgsz=640)
        boxes_list = results[0].boxes.data.tolist()
        if boxes_list is None or len(boxes_list) == 0:
            print("没有检测到任何物体")
        print("detect ===========begin=============")
        print(boxes_list)
        print("detect ===========end=============")
        target_box_list = []
        for row in boxes_list:
            scale = float(row[4])
            scale = round(scale, 2)
            id = row[5]
            if id == int(sku_id) and scale >= float(sku_scale) :
                x1,y1,x2,y2 = row[0],row[1],row[2],row[3]
                target_c_x = int((x2-x1)/2+x1)
                target_c_y = int((y2-y1)/2+y1)
                center_coordinates = (target_c_x, target_c_y)
                radius = 5  
                color = (0, 0, 255)  # B=0, G=0, R=255  
                image = cv2.circle(latest_color_frame, center_coordinates, radius, color, -1)
                # 获取深度值
                depth_value = latest_depth_frame[target_c_y-1, target_c_x-1] 
                depth_value1 = latest_depth_frame[target_c_y-1, target_c_x] 
                depth_value2 = latest_depth_frame[target_c_y-1, target_c_x+1] 
                depth_value3 = latest_depth_frame[target_c_y, target_c_x-1] 
                depth_value4 = latest_depth_frame[target_c_y, target_c_x] 
                depth_value5 = latest_depth_frame[target_c_y, target_c_x+1] 
                depth_value6 = latest_depth_frame[target_c_y+1, target_c_x-1] 
                depth_value7 = latest_depth_frame[target_c_y+1, target_c_x] 
                depth_value8 = latest_depth_frame[target_c_y+1, target_c_x+1] 
                depth_list = [depth_value,depth_value1,depth_value2,depth_value3,depth_value4,depth_value5,depth_value6,depth_value7,depth_value8]
                depth_list = filter_invalid_data(list(depth_list))
                if len(depth_list) == 0:
                    print("没有有效深度值")
                    continue
                average_depth = np.mean(list(depth_list)) 
                # depth_value = round(depth_value, 2)
                print(f"像素坐标: ({target_c_x}, {target_c_y})  深度值: {average_depth} mm")
                cam_xyz = pixel2point3d(target_c_x, target_c_y, average_depth, intr)
                python_floats = [round(float(x),3) for x in cam_xyz] 
                print(f"相机坐标系下的三维坐标: {python_floats}")

                top_left = (int(x1), int(y1))
                bottom_right = (int(x2), int(y2))
                color = (0, 255, 0)  # BGR格式的颜色，绿色
                cv2.rectangle(latest_color_frame, top_left, bottom_right, color, thickness=2)
                label = str(id)+":"+str(scale)+":"+str(depth_value)
                cv2.putText(latest_color_frame, label, (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
                target_box_list.append([int(x1),int(y1),int(x2),int(y2),int(average_depth)])
        mindepth_box = find_nearest_rectangle(target_box_list)
        print("距离最近的box")
        print(mindepth_box)
        draw_top1(latest_color_frame,mindepth_box)
        box_c_x = int((mindepth_box[2]-mindepth_box[0])/2+mindepth_box[0])
        box_c_y = int((mindepth_box[3]-mindepth_box[1])/2+mindepth_box[1])
        cam_xyz = pixel2point3d(box_c_x, box_c_y, mindepth_box[4], intr)
        target_xyz = [round(float(x),3) for x in cam_xyz] 
        print(f"相机坐标系下的三维坐标: {target_xyz}")
        # 保存临时图片
        timestamp = int(time.time())
        filename = f"static/capture_{timestamp}.jpg"
        cv2.imwrite(filename, latest_color_frame)

        return jsonify({
            "image_url": filename,
            "message": "查找成功",
            "mindepth_box":mindepth_box,
            "target_xyz":target_xyz
        })


@app.route('/detect')
def detect():
    with frame_lock:
        if latest_color_frame is None:
            return jsonify({"error": "相机未就绪"}), 500
            
        # 保存临时图片
        timestamp = int(time.time())
        filename = f"static/capture_{timestamp}.jpg"
        cv2.imwrite(filename, latest_color_frame)
        results = model.predict(latest_color_frame, conf=0.5, imgsz=640)
        boxes_list = results[0].boxes.data.tolist()
        if boxes_list is None or len(boxes_list) == 0:
            print("没有检测到任何物体")
        print("detect ===========begin=============")
        print(boxes_list)
        print("detect ===========end=============")
        for row in boxes_list:
            x1,y1,x2,y2 = row[0],row[1],row[2],row[3]
            scale = float(row[4])
            scale = round(scale, 2)
            id = row[5]

            target_c_x = int((x2-x1)/2+x1)
            target_c_y = int((y2-y1)/2+y1)
            center_coordinates = (target_c_x, target_c_y)
            radius = 5  
            color = (0, 0, 255)  # B=0, G=0, R=255  
            image = cv2.circle(latest_color_frame, center_coordinates, radius, color, -1)
            # 获取深度值
            depth_value = latest_depth_frame[target_c_y, target_c_x] 
            # depth_value = round(depth_value, 2)
            print(f"像素坐标: ({target_c_x}, {target_c_y})  深度值: {depth_value} mm")
            cam_xyz = pixel2point3d(target_c_x, target_c_y, depth_value, intr)
            python_floats = [round(float(x),3) for x in cam_xyz] 
            print(f"相机坐标系下的三维坐标: {cam_xyz}")
            print(python_floats)

            top_left = (int(x1), int(y1))
            bottom_right = (int(x2), int(y2))
            color = (0, 255, 0)  # BGR格式的颜色，绿色
            cv2.rectangle(latest_color_frame, top_left, bottom_right, color, thickness=2)
            label = str(id)+":"+str(scale)+":"+str(depth_value)
            cv2.putText(latest_color_frame, label, (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
        # 保存临时图片
        timestamp = int(time.time())
        filename = f"static/capture_{timestamp}.jpg"
        cv2.imwrite(filename, latest_color_frame)

        return jsonify({
            "image_url": filename,
            "message": "检测成功",
            "boxes_list":boxes_list
        })
    

@app.route('/capture')
def capture():
    """拍照接口"""
    with frame_lock:
        if latest_color_frame is None:
            return jsonify({"error": "相机未就绪"}), 500
            
        # 保存临时图片
        timestamp = int(time.time())
        filename = f"static/capture_{timestamp}.jpg"
        cv2.imwrite(filename, latest_color_frame)
        
        return jsonify({
            "image_url": filename,
            "message": "拍照成功"
        })

def generate():
    """生成实时视频流"""
    while True:
        with frame_lock:
            if latest_color_frame is None:
                continue
                
            # 编码为JPEG
            ret, jpeg = cv2.imencode('.jpg', latest_color_frame)
            frame = jpeg.tobytes()
            
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')

@app.route('/video_feed')
def video_feed():
    """视频流接口"""
    return Response(generate(),
                   mimetype='multipart/x-mixed-replace; boundary=frame')

if __name__ == '__main__':
    # 启动相机线程
    threading.Thread(target=camera_thread, daemon=True).start()
    
    # 创建静态文件目录
    import os
    os.makedirs('static', exist_ok=True)
    
    app.run(host='0.0.0.0', port=5000, threaded=True)